#include "KVStorageClient.h"
#include <iostream>
#include <fstream>
#include <sys/mman.h>
#include <string.h>
#include <gtest/gtest.h>
#include <vector>
#include "Logger.h"

using namespace std;

const char *SHM_PATH = "/dev/shm/kvclient_shm";
const char *SHM_FILE = "kvclient_shm";
const char *MNT_PATH = "/mnt/kvclient/";
const uint64_t SHM_SIZE = 40960;
uint64_t segOff = 4096;
uint64_t segSize = 4096;
uint64_t tokenOff = 1024;
uint64_t tokenLen = 1024;
int64_t objId = 4020;

class StorageTest : public ::testing::Test {
protected:
    void SetUp() override
    {
        for (const auto &entry : fs::directory_iterator(".")) {
            if (entry.path().string().find("test_log") != std::string::npos) {
                fs::remove(entry.path());
            }
        }

        Logger::instance().init("test_log", 1000);
    }

    void TearDown() override
    {
        Logger::instance().shutdown();

        for (const auto &entry : fs::directory_iterator(".")) {
            if (entry.path().string().find("test_log") != std::string::npos) {
                fs::remove(entry.path());
            }
        }
    }

public:
    // 准备共享内存
    void prepareSharedMemory(const char *file, const uint64_t size, uint8_t **mapAddr)
    {
        int32_t fd = shm_open(file, O_CREAT | O_RDWR, 0644);
        ASSERT_NE(-1, fd);

        int32_t ret = ftruncate(fd, size);
        ASSERT_EQ(0, ret);

        *mapAddr = (uint8_t *)mmap(nullptr, size, PROT_READ | PROT_WRITE, MAP_SHARED, fd, 0);
        ASSERT_NE(MAP_FAILED, *mapAddr);
        ASSERT_NE(nullptr, *mapAddr);

        std::fill(*mapAddr, (*mapAddr) + size, '0');
        close(fd);
    }
    void checkSharedMemory(uint8_t *mapAddr, const uint64_t off, const uint64_t len, char tgt)
    {
        for (uint64_t i = 0; i < len; ++i) {
            if (mapAddr[i + off] != tgt) {
                std::cout << "first mismatch idx(" << i + off << "):";
                EXPECT_EQ(tgt, mapAddr[i]);
                return;
            }
        }
    }
};

TEST_F(StorageTest, PosixPutGetRemove)
{
    std::string objIdStr = std::to_string(objId) + "_" + "1";
    uint8_t *mapAddr = nullptr;
    prepareSharedMemory(SHM_FILE, SHM_SIZE, &mapAddr);

    KVStorageClient &client = KVStorageClient::getInstance();
    client.init(FileSystem::POSIX, SHM_PATH, mapAddr, MNT_PATH);
    LOG_INFO_FMT("init success");
    client.remove(objIdStr);

    std::fill(mapAddr + segOff + tokenOff, mapAddr + segOff + tokenOff + tokenLen, 'a');
    Segment segment(objId, segOff, mapAddr + segOff, segSize, false, SHM_PATH);
    EXPECT_EQ(0, client.put(objIdStr, segment, tokenOff, tokenOff, tokenLen));
    std::fill(mapAddr + segOff + tokenOff + tokenLen, mapAddr + segOff + tokenOff + tokenLen + tokenLen, 'b');
    EXPECT_EQ(0, client.put(objIdStr, segment, tokenOff + tokenLen, tokenOff + tokenLen, tokenLen));

    // 清空mapAddr再读取
    std::fill(mapAddr, mapAddr + SHM_SIZE, '0');
    EXPECT_EQ(0, client.get(objIdStr, segment, tokenOff, tokenOff, tokenLen));
    EXPECT_EQ(0, client.get(objIdStr, segment, tokenOff + tokenLen, tokenOff + tokenLen, tokenLen));

    checkSharedMemory(mapAddr, segOff + tokenOff, tokenLen, 'a');

    EXPECT_EQ(0, munmap(mapAddr, SHM_SIZE));
}

TEST_F(StorageTest, CheckSubDirAndCread)
{
    int32_t file_count = 0;
    uint8_t *mapAddr = nullptr;
    prepareSharedMemory(SHM_FILE, SHM_SIZE, &mapAddr);

    KVStorageClient &client = KVStorageClient::getInstance();
    client.init(FileSystem::POSIX, SHM_PATH, mapAddr, MNT_PATH);
    LOG_INFO_FMT("init success");
    
    int64_t objId1 = 4096;
    std::string objId1Str = std::to_string(objId1) + "_" + "1";
    int64_t objId2 = 4096 + 1;
    std::string objId2Str = std::to_string(objId2) + "_" + "1";
    int64_t objId3 = 8192;
    std::string objId3Str = std::to_string(objId3) + "_" + "1";
    client.remove(objId1Str);
    client.remove(objId2Str);
    client.remove(objId3Str);

    std::fill(mapAddr, mapAddr + tokenLen, 'a');
    Segment segment1(objId1, 0, mapAddr, segSize, false, SHM_PATH);
    Segment segment2(objId2, segSize, mapAddr, segSize, false, SHM_PATH);
    Segment segment3(objId3, segSize * 2, mapAddr, segSize, false, SHM_PATH);

    // 检查子目录创建
    EXPECT_EQ(0, client.put(objId1Str, segment1, 0, 0, tokenLen));
    EXPECT_EQ(true, fs::exists(std::string(MNT_PATH) + "/NHD/" + std::to_string(objId1 % INDEX_MOD)) && fs::is_directory(std::string(MNT_PATH) + "/NHD/" + std::to_string(objId1 % INDEX_MOD)));
    EXPECT_EQ(true, fs::exists(std::string(MNT_PATH) + "/NHD/" + std::to_string(objId1 % INDEX_MOD) + "/kvcache_4096_1"));
    
    // 确认打散到不同文件夹
    EXPECT_EQ(0, client.put(objId2Str, segment2, 0, segSize, tokenLen));
    EXPECT_EQ(true, fs::exists(std::string(MNT_PATH) + "/NHD/" + std::to_string(objId2 % INDEX_MOD)) && fs::is_directory(std::string(MNT_PATH) + "/NHD/" + std::to_string(objId2 % INDEX_MOD)));
    EXPECT_EQ(true, fs::exists(std::string(MNT_PATH) + "/NHD/" + std::to_string(objId2 % INDEX_MOD) + "/kvcache_4097_1"));
    
    // 测试取模相同场景
    EXPECT_EQ(0, client.put(objId3Str, segment3, 0, segSize * 2, tokenLen));
    EXPECT_EQ(true, fs::exists(std::string(MNT_PATH) + "/NHD/" + std::to_string(objId3 % INDEX_MOD)) && fs::is_directory(std::string(MNT_PATH) + "/NHD/" + std::to_string(objId3 % INDEX_MOD)));
    EXPECT_EQ(true, fs::exists(std::string(MNT_PATH) + "/NHD/" + std::to_string(objId3 % INDEX_MOD) + "/kvcache_8192_1"));
    for (const auto& entry : fs::directory_iterator(std::string(MNT_PATH) + "/NHD/" + std::to_string(objId3 % INDEX_MOD))) {
        if (entry.is_regular_file()) {
            file_count++;
        }
    }
    EXPECT_EQ(2, file_count);
    EXPECT_EQ(0, munmap(mapAddr, SHM_SIZE));
}

int main(int argc, char **argv)
{
    testing::InitGoogleTest(&argc, argv);
    return RUN_ALL_TESTS();
}